from langchain_ollama import OllamaLLM
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain.memory import SimpleMemory
from langchain.chains import LLMChain

# 1. 初始化本地 Ollama 模型
ollama_llm = OllamaLLM(
    model="llama3.1"
)

# 2. 带上下文的提示语
template_with_history = """你是一位贴心的日常生活助理，以下是我们的对话历史：
{history}
现在根据最后一句话来回答用户问题：
{human_input}
"""

prompt_with_history = PromptTemplate.from_template(template_with_history)

parser = StrOutputParser()

# 3. 初始化memory
memory = SimpleMemory()

# 4. 创建支持记忆的 chain
# 注意这里我们使用 RunnableSequence 替代了 LLMChain
memory_chains = LLMChain(
    llm=ollama_llm,
    prompt=prompt_with_history,  # 正确命名 prompt 参数
    output_parser=parser,  # 确保这是被接受的参数
)

# 5. 多轮对话测试
user_input1 = "你好啊，我今天早上吃了火腿蛋面包，有什么搭配饮料的推荐吗？"
response1 = memory_chains.predict({"human_input": user_input1})  # 确保使用 run 方法
print("回答1：", response1)

user_input2 = "我下午打算去健身，能再推荐一下轻食吗？"
response2 = memory_chains.predict({"human_input": user_input2})
print("回答2：", response2)